from cryptography.hazmat.primitives import serialization, hashes, padding
from cryptography.hazmat.primitives.asymmetric import rsa, padding as asym_padding
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
import os
import base64

class CryptoUtils:
    @staticmethod
    def generate_rsa_keys():
        """生成RSA密钥对"""
        private_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048,
            backend=default_backend()
        )
        public_key = private_key.public_key()
        
        # 序列化私钥
        private_pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        )
        
        # 序列化公钥
        public_pem = public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        
        return private_pem, public_pem

    @staticmethod
    def rsa_encrypt(public_key_pem, message):
        """使用RSA公钥加密消息"""
        public_key = serialization.load_pem_public_key(
            public_key_pem,
            backend=default_backend()
        )
        
        # 确保消息是字节类型
        if isinstance(message, str):
            message = message.encode('utf-8')
        
        encrypted = public_key.encrypt(
            message,
            asym_padding.OAEP(
                mgf=asym_padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        return base64.b64encode(encrypted)

    @staticmethod
    def rsa_decrypt(private_key_pem, encrypted_message):
        """使用RSA私钥解密消息"""
        private_key = serialization.load_pem_private_key(
            private_key_pem,
            password=None,
            backend=default_backend()
        )
        
        # 确保加密消息是字节类型
        if isinstance(encrypted_message, str):
            encrypted_message = encrypted_message.encode('utf-8')
        
        encrypted = base64.b64decode(encrypted_message)
        decrypted = private_key.decrypt(
            encrypted,
            asym_padding.OAEP(
                mgf=asym_padding.MGF1(algorithm=hashes.SHA256()),
                algorithm=hashes.SHA256(),
                label=None
            )
        )
        return decrypted.decode('utf-8')

    @staticmethod
    def generate_aes_key():
        """生成AES密钥"""
        return os.urandom(32)

    @staticmethod
    def aes_encrypt(key, message):
        """使用AES加密消息"""
        # 确保消息是字符串类型
        if not isinstance(message, str):
            message = str(message)
            
        iv = os.urandom(16)
        cipher = Cipher(
            algorithms.AES(key),
            modes.CBC(iv),
            backend=default_backend()
        )
        encryptor = cipher.encryptor()
        
        # 填充消息
        padder = padding.PKCS7(algorithms.AES.block_size).padder()
        padded_data = padder.update(message.encode('utf-8')) + padder.finalize()
        
        encrypted = encryptor.update(padded_data) + encryptor.finalize()
        return base64.b64encode(iv + encrypted)

    @staticmethod
    def aes_decrypt(key, encrypted_message):
        """使用AES解密消息"""
        # 确保加密消息是字节类型
        if isinstance(encrypted_message, str):
            encrypted_message = encrypted_message.encode('utf-8')
            
        encrypted = base64.b64decode(encrypted_message)
        iv = encrypted[:16]
        ciphertext = encrypted[16:]
        
        cipher = Cipher(
            algorithms.AES(key),
            modes.CBC(iv),
            backend=default_backend()
        )
        decryptor = cipher.decryptor()
        
        padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize()
        unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
        plaintext = unpadder.update(padded_plaintext) + unpadder.finalize()
        
        return plaintext.decode('utf-8')